

'''
description: 
first benchmark IoU: 0.5, 0.75 for AP 
second benchmark IoU: 0.5:0.95, step = 0.05
return {*}
#*  
'''


import os 
import sys
os.chdir("/data/xusc/exp/topictrack-bee")


from os.path import join, isdir, isfile, exists
from os import makedirs

import cv2
from glob import glob
import json

from os.path import  join, split, exists
import numpy as np
import glob

from tools.det_acc import  *


class DetAccCalculator:

    def __init__(self, exp_name  ):

        
        gt_root = "/data/xusc/exp/topictrack-bee/results/gt"
        exp_result_root = "/data/xusc/exp/topictrack-bee/results/trackers"


        #* load prediction 

        seqs = {}
        for seq_path in glob.glob(join(gt_root,exp_name,"bee*")):
            # records,name_dict = read_record(join(seq_path,'gt','gt.txt'))
            seq_name = seq_path.split('/')[-1]
            gt_s = np.loadtxt(join(seq_path,'gt','gt.txt'),dtype=np.float32,delimiter=',')
            # gt_s = np.loadtxt(join(seq_path,'gt','gt.txt'),dtype=np.str0,delimiter=',')

            n = gt_s.shape[0]
            new_gt = np.concatenate([gt_s[:,0][:,None],np.zeros([n,1]),np.ones([n,1]),gt_s[:,2][:,None], gt_s[:,3][:,None],(gt_s[:,2] +  gt_s[:,4])[:,None] , (gt_s[:,3] +  gt_s[:,5])[:,None] ],axis=1).astype(np.int32)
            
            seqs[seq_name] = new_gt


        pred_seqs = {}
        mode = '_test_post'    
        exp_name_prefix = exp_name.split('-')[0]

        for seq_path in glob.glob(join(exp_result_root,exp_name, exp_name_prefix+ mode, 'data',"bee*filter*")):
            # records,name_dict = read_record(join(seq_path,'gt','gt.txt'))
            seq_name = seq_path.split('/')[-1].split('.')[0]
            gt_s = np.loadtxt(seq_path,dtype=np.float32,delimiter=',')

            n = gt_s.shape[0]
            new_gt = np.concatenate([gt_s[:,0][:,None],np.zeros([n,1]),np.ones([n,1]),gt_s[:,2][:,None], gt_s[:,3][:,None],(gt_s[:,2] +  gt_s[:,4])[:,None] , (gt_s[:,3] +  gt_s[:,5])[:,None] ],axis=1).astype(np.int32)
            pred_seqs[seq_name] = new_gt

        self.gts = seqs
        self.preds = pred_seqs
        # assert  len(self.gts.keys()) == len(self.preds.keys())

        self.bee_names = list(self.preds.keys())


    def __len__(self):
        return len(self.bee_names )

    def getitem(self,idx):
        # return self.gts[self.bee_names[idx]],self.preds[self.bee_names[idx]]
        return self.gts[self.bee_names[idx].split('_')[0]],self.preds[self.bee_names[idx]],


    def get_ap(self, idx,  iou = 0.5):
        
        print('name:', self.bee_names[idx])

        gt, predicted = self.getitem(idx)

        ap_all = []
        for i in range(1,gt[-1][0]+1):
            # print('id ',i)
            gt_one = [row for row in gt if row[0] == i]
            pre_one = [row for row in predicted if row[0] == i]
            ap = compute_ap(gt_one,pre_one,iou)
            ap_all.append(ap)

        return sum(ap_all)/gt[-1][0]

    



class AntDetAccCalculator(DetAccCalculator):

    def __init__(self, exp_name  ):
        
        gt_root = "/data/xusc/exp/topictrack-bee/results/gt"
        exp_result_root = "/data/xusc/exp/topictrack-bee/results/trackers"


        #* load prediction 
        corrected_sequence_name = 'antmove-val'
        seqs = {}
        for seq_path in glob.glob(join(gt_root,corrected_sequence_name,"ant*")):
            # records,name_dict = read_record(join(seq_path,'gt','gt.txt'))
            seq_name = seq_path.split('/')[-1]
            gt_s = np.loadtxt(join(seq_path,'gt','gt.txt'),dtype=np.float32,delimiter=',')
            # gt_s = np.loadtxt(join(seq_path,'gt','gt.txt'),dtype=np.str0,delimiter=',')

            n = gt_s.shape[0]
            new_gt = np.concatenate([gt_s[:,0][:,None],np.zeros([n,1]),np.ones([n,1]),gt_s[:,2][:,None], gt_s[:,3][:,None],(gt_s[:,2] +  gt_s[:,4])[:,None] , (gt_s[:,3] +  gt_s[:,5])[:,None] ],axis=1).astype(np.int32)
            
            seqs[seq_name] = new_gt


        pred_seqs = {}
        mode = '_test_post'    
        exp_name_prefix = exp_name.split('-')[0]

        for seq_path in glob.glob(join(exp_result_root,exp_name, exp_name_prefix+ mode, 'data',"*ant*")):
            # records,name_dict = read_record(join(seq_path,'gt','gt.txt'))
            seq_name = seq_path.split('/')[-1].split('.')[0]
            gt_s = np.loadtxt(seq_path,dtype=np.float32,delimiter=',')

            n = gt_s.shape[0]
            new_gt = np.concatenate([gt_s[:,0][:,None],np.zeros([n,1]),np.ones([n,1]),gt_s[:,2][:,None], gt_s[:,3][:,None],(gt_s[:,2] +  gt_s[:,4])[:,None] , (gt_s[:,3] +  gt_s[:,5])[:,None] ],axis=1).astype(np.int32)
            pred_seqs[seq_name] = new_gt

        self.gts = seqs
        self.preds = pred_seqs
        # assert  len(self.gts.keys()) == len(self.preds.keys())
        self.bee_names = list(self.preds.keys())

    def get_ap(self, idx,  iou = 0.5):
        
        print('name:', self.bee_names[idx])

        gt, predicted = self.getitem(idx)

        ap_all = []
        for i in range(1,gt[-1][0]+1):
            # print('id ',i)
            gt_one = [row for row in gt if row[0] == i]
            pre_one = [row for row in predicted if row[0] == i]
            ap = compute_ap(gt_one,pre_one,iou)
            ap_all.append(ap)

        return sum(ap_all)/gt[-1][0]


        
if __name__ =="__main__":
    det_acc_calculator = DetAccCalculator('beedance-val')
    for idx in range(det_acc_calculator.__len__()):
        print(det_acc_calculator.get_ap(idx))


